package de.h_da.fbi.cg.pointcloud.spatial;

import java.lang.reflect.Array;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.Set;

/**
 * <p>
 * A k-d tree divides a k-dimensional space relative to the points it contains
 * by storing them in a binary tree, discriminating by a different dimension at
 * each level of the tree. It allows efficient point data retrieval (
 * <em>O(lg(n))</em>) and range searching.
 * </p>
 * 
 * <p>
 * KDTree conforms to the java.util.Map interface except that Iterator.remove is
 * not supported by the returned views.
 * </p>
 */
public class KDTree<C extends Comparable<C>, P extends IPoint<C>, V> implements IRangeSearchTree<C, P, V> {
	// ===========================================================
	// Constants
	// ===========================================================

	// ===========================================================
	// Fields
	// ===========================================================

	int mSize;
	int mHashCode;
	int mDimensions;
	KDNode mRoot;

	// ===========================================================
	// Constructors
	// ===========================================================

	/**
	 * Creates a two-dimensional KDTree.
	 */
	public KDTree() {
		this(2);
	}

	/**
	 * Creates a KDTree of the specified number of dimensions.
	 * 
	 * @param dimensions
	 *            The number of dimensions. Must be greater than 0.
	 */
	public KDTree(final int dimensions) {
		assert (dimensions > 0);
		this.mDimensions = dimensions;
		this.clear();
	}

	// ===========================================================
	// Getter & Setter
	// ===========================================================

	// ===========================================================
	// Methods for/from SuperClass/Interfaces
	// ===========================================================

	@Override
	public void clear() {
		this.mRoot = null;
		this.mSize = this.mHashCode = 0;
	}

	@Override
	@SuppressWarnings("unchecked")
	public boolean containsKey(final Object pKey) throws ClassCastException {
		return this.getNode((P) pKey) != null;
	}

	@Override
	public boolean containsValue(final Object pValue) {
		return this.findValue(this.mRoot, pValue) != null;
	}

	@Override
	public Set<Map.Entry<P, V>> entrySet() {
		return new MapEntrySet();
	}

	@Override
	public boolean equals(final Object pOther) throws ClassCastException {
		if(!(pOther instanceof Map<?, ?>)) {
			return false;
		}

		if(pOther == this) {
			return true;
		}

		final Map<?, ?> map = (Map<?, ?>) pOther;

		return this.entrySet().equals(map.entrySet());
	}

	@Override
	@SuppressWarnings("unchecked")
	public V get(final Object point) throws ClassCastException {
		final KDNode node = this.getNode((P) point);

		return (node == null ? null : node.getValue());
	}

	@Override
	public int hashCode() {
		return this.mHashCode;
	}

	@Override
	public boolean isEmpty() {
		return (this.mRoot == null);
	}

	@Override
	public Set<P> keySet() {
		return new KeySet();
	}

	/**
	 * @return The old value if an existing value is replaced by the inserted value.
	 */
	@Override
	@SuppressWarnings("unchecked")
	public V put(final P point, final V value) {
		final KDNode[] parent = new KDTree.KDNode[1];
		KDNode node = this.getNode(point, parent);
		V old = null;

		if(node != null) {
			old = node.getValue();
			this.mHashCode -= node.hashCode();
			node.mValue = value;
		} else {
			if(parent[0] == null) {
				node = this.mRoot = new KDNode(0, point, value);
			} else {
				final int dimension = parent[0].mDimension;

				if(point.getCoordinate(dimension).compareTo(parent[0].mPoint.getCoordinate(dimension)) >= 0) {
					node = parent[0].mHigher = new KDNode((dimension + 1) % this.mDimensions, point, value);
				} else {
					node = parent[0].mLower = new KDNode((dimension + 1) % this.mDimensions, point, value);
				}
			}

			this.mSize++;
		}

		this.mHashCode += node.hashCode();

		return old;
	}

	@Override
	public void putAll(final Map<? extends P, ? extends V> map) {
		for(final Map.Entry<? extends P, ? extends V> pair : map.entrySet()) {
			this.put(pair.getKey(), pair.getValue());
		}
	}

	@Override
	@SuppressWarnings("unchecked")
	public V remove(final Object key) throws ClassCastException {
		final KDNode[] parent = new KDTree.KDNode[1];
		KDNode node = this.getNode((P) key, parent);
		V old = null;

		if(node != null) {
			final KDNode child = node;

			node = this.recursiveRemoveNode(child);

			if(parent[0] == null) {
				this.mRoot = node;
			} else if(child == parent[0].mLower) {
				parent[0].mLower = node;
			} else if(child == parent[0].mHigher) {
				parent[0].mHigher = node;
			}

			--this.mSize;
			this.mHashCode -= child.hashCode();
			old = child.getValue();
		}

		return old;
	}

	@Override
	public int size() {
		return this.mSize;
	}

	@Override
	public Collection<V> values() {
		return new ValueCollection();
	}

	@Override
	public Iterator<Map.Entry<P, V>> iterator(final P lower, final P upper) {
		return new MapEntryIterator(lower, upper);
	}

	// ===========================================================
	// Methods
	// ===========================================================

	// ===========================================================
	// Inner and Anonymous Classes
	// ===========================================================

	KDNode getNode(final P pPoint, final KDNode[] pParent) {
		int dimension;
		KDNode node = this.mRoot;
		KDNode current;
		KDNode last = null;
		C c1;
		C c2;

		while(node != null) {
			dimension = node.mDimension;
			c1 = pPoint.getCoordinate(dimension);
			c2 = node.mPoint.getCoordinate(dimension);
			current = node;

			if(c1.compareTo(c2) > 0) {
				node = node.mHigher;
			} else if(c1.compareTo(c2) < 0) {
				node = node.mLower;
			} else if(node.mPoint.equals(pPoint)) {
				if(pParent != null) {
					pParent[0] = last;
				}
				return node;
			} else {
				node = node.mHigher;
			}

			last = current;
		}

		if(pParent != null) {
			pParent[0] = last;
		}

		return null;
	}

	KDNode getNode(final P pPoint) {
		return this.getNode(pPoint, null);
	}

	@SuppressWarnings("unchecked")
	KDNode getMinimumNode(final KDNode pKDNodeCompare, final KDNode pKDNode, final int pDimension, final KDNode[] pParentKDNode) {
		KDNode result;

		if(pDimension == pKDNodeCompare.mDimension) {
			if(pKDNodeCompare.mLower != null) {
				return this.getMinimumNode(pKDNodeCompare.mLower, pKDNodeCompare, pDimension, pParentKDNode);
			} else {
				result = pKDNodeCompare;
			}
		} else {
			KDNode nlow = null, nhigh = null;
			final KDNode[] plow = new KDTree.KDNode[1];
			final KDNode[] phigh = new KDTree.KDNode[1];

			if(pKDNodeCompare.mLower != null) {
				nlow = this.getMinimumNode(pKDNodeCompare.mLower, pKDNodeCompare, pDimension, plow);
			}

			if(pKDNodeCompare.mHigher != null) {
				nhigh = this.getMinimumNode(pKDNodeCompare.mHigher, pKDNodeCompare, pDimension, phigh);
			}

			if(nlow != null && nhigh != null) {
				if(nlow.mPoint.getCoordinate(pDimension).compareTo(nhigh.mPoint.getCoordinate(pDimension)) < 0) {
					result = nlow;
					pParentKDNode[0] = plow[0];
				} else {
					result = nhigh;
					pParentKDNode[0] = phigh[0];
				}
			} else if(nlow != null) {
				result = nlow;
				pParentKDNode[0] = plow[0];
			} else if(nhigh != null) {
				result = nhigh;
				pParentKDNode[0] = phigh[0];
			} else {
				result = pKDNodeCompare;
			}
		}

		if(result == pKDNodeCompare) {
			pParentKDNode[0] = pKDNode;
		} else if(pKDNodeCompare.mPoint.getCoordinate(pDimension).compareTo(result.mPoint.getCoordinate(pDimension)) < 0) {
			result = pKDNodeCompare;
			pParentKDNode[0] = pKDNode;
		}

		return result;
	}

	@SuppressWarnings("unchecked")
	KDNode recursiveRemoveNode(final KDNode pKDNode) {
		int dimension;

		if(pKDNode.mLower == null && pKDNode.mHigher == null) {
			return null;
		} else {
			dimension = pKDNode.mDimension;
		}

		if(pKDNode.mHigher == null) {
			pKDNode.mHigher = pKDNode.mLower;
			pKDNode.mLower = null;
		}

		final KDNode[] parent = new KDTree.KDNode[1];
		final KDNode newRoot = this.getMinimumNode(pKDNode.mHigher, pKDNode, dimension, parent);
		final KDNode child = this.recursiveRemoveNode(newRoot);

		if(parent[0].mLower == newRoot) {
			parent[0].mLower = child;
		} else {
			parent[0].mHigher = child;
		}

		newRoot.mLower = pKDNode.mLower;
		newRoot.mHigher = pKDNode.mHigher;
		newRoot.mDimension = pKDNode.mDimension;

		return newRoot;
	}

	KDNode findValue(final KDNode node, final Object value) {
		if(node == null || (value == null ? node.getValue() == null : value.equals(node.getValue()))) {
			return node;
		}

		KDNode result;

		if((result = this.findValue(node.mLower, value)) == null) {
			result = this.findValue(node.mHigher, value);
		}

		return result;
	}

	boolean isInRange(final P pPoint, final P pLower, final P pHigher) {
		C coordinate1, coordinate2 = null, coordinate3 = null;

		if(pLower != null || pHigher != null) {
			int dimensions;
			dimensions = pPoint.getDimensions();

			for(int i = 0; i < dimensions; ++i) {
				coordinate1 = pPoint.getCoordinate(i);
				if(pLower != null) {
					coordinate2 = pLower.getCoordinate(i);
				}
				if(pHigher != null) {
					coordinate3 = pHigher.getCoordinate(i);
				}
				if((coordinate2 != null && coordinate1.compareTo(coordinate2) < 0) || (coordinate3 != null && coordinate1.compareTo(coordinate3) > 0)) {
					return false;
				}
			}
		}

		return true;
	}


	int fillArray(final KDNode[] pKDNodes, int pIndex, final KDNode pKDNode) {
		if(pKDNode == null) {
			return pIndex;
		}
		pKDNodes[pIndex] = pKDNode;
		pIndex = this.fillArray(pKDNodes, pIndex + 1, pKDNode.mLower);
		return this.fillArray(pKDNodes, pIndex, pKDNode.mHigher);
	}

	public void balance() {
		if(this.isEmpty()) {
			return;
		}

		final KDNode[] nodes = (KDNode[]) Array.newInstance(KDNode.class, this.size());
		this.fillArray(nodes, 0, this.mRoot);

		this.mRoot = this.balance(nodes, 0, nodes.length, new KDNodeComparator());
	}

	private KDNode balance(final KDNode[] nodes, final int pStart, final int pEnd, final KDNodeComparator pKDNodeComparator) {
		KDNode midNode = null;
		final int size = pEnd - pStart;

		if(size > 1) {
			int nth = pStart + (size >> 1);
			int nthPrevious = nth - 1;
			int dimension = pKDNodeComparator.getDimension();

			Arrays.sort(nodes, pStart, pEnd, pKDNodeComparator);

			while(nth > pStart && nodes[nth].mPoint.getCoordinate(dimension).compareTo(nodes[nthPrevious].mPoint.getCoordinate(dimension)) == 0) {
				--nth;
				--nthPrevious;
			}

			midNode = nodes[nth];
			midNode.mDimension = dimension;

			if(++dimension >= this.mDimensions) {
				dimension = 0;
			}

			pKDNodeComparator.setDimension(dimension);

			midNode.mLower = this.balance(nodes, pStart, nth, pKDNodeComparator);

			pKDNodeComparator.setDimension(dimension);

			midNode.mHigher = this.balance(nodes, nth + 1, pEnd, pKDNodeComparator);
		} else if(size == 1) {
			midNode = nodes[pStart];
			midNode.mDimension = pKDNodeComparator.getDimension();
			midNode.mLower = midNode.mHigher = null;
		}

		return midNode;
	}

	final class KDNodeComparator implements Comparator<KDNode> {
		int mDimension = 0;

		void setDimension(final int pDimension) {
			this.mDimension = pDimension;
		}

		int getDimension() {
			return this.mDimension;
		}

		@Override
		public int compare(final KDNode pKDNodeA, final KDNode pKDNodeB) {
			return pKDNodeA.mPoint.getCoordinate(this.mDimension).compareTo(pKDNodeB.mPoint.getCoordinate(this.mDimension));
		}
	}

	final class KDNode implements Map.Entry<P, V> {
		int mDimension;
		P mPoint;
		V mValue;
		KDNode mLower;
		KDNode mHigher;

		KDNode(final int pDimension, final P pPoint, final V pValue) {
			this.mPoint = pPoint;
			this.mValue = pValue;
			this.mLower = this.mHigher = null;
			this.mDimension = pDimension;
		}

		@Override
		public boolean equals(final Object pOther) {
			final KDNode node = (KDNode) pOther;

			if(node == this) {
				return true;
			}

			return ((this.getKey() == null ? node.getKey() == null : this.getKey().equals(node.getKey())) && (this.getValue() == null ? node.getValue() == null : this.getValue().equals(node.getValue())));
		}

		@Override
		public P getKey() {
			return this.mPoint;
		}

		@Override
		public V getValue() {
			return this.mValue;
		}

		@Override
		public V setValue(final V pValue) {
			final V old = this.mValue;
			KDTree.this.mHashCode -= this.hashCode();
			this.mValue = pValue;
			KDTree.this.mHashCode += this.hashCode();
			return old;
		}

		@Override
		public int hashCode() {
			return ((this.getKey() == null ? 0 : this.getKey().hashCode()) ^ (this.getValue() == null ? 0 : this.getValue().hashCode()));
		}
	}


	final class MapEntryIterator implements Iterator<Map.Entry<P, V>> {
		LinkedList<KDNode> mKDNodeStack;
		KDNode mNextKDNode;
		P mLower, mHigher;

		MapEntryIterator(final P pLower, final P pHigher) {
			this.mKDNodeStack = new LinkedList<KDNode>();
			this.mLower = pLower;
			this.mHigher = pHigher;
			this.mNextKDNode = null;

			if(KDTree.this.mRoot != null) {
				this.mKDNodeStack.addLast(KDTree.this.mRoot);
			}
			this.next();
		}

		MapEntryIterator() {
			this(null, null);
		}

		@Override
		public boolean hasNext() {
			return (this.mNextKDNode != null);
		}

		@Override
		public Map.Entry<P, V> next() {
			final KDNode old = this.mNextKDNode;

			while(!this.mKDNodeStack.isEmpty()) {
				final KDNode node = this.mKDNodeStack.removeLast();
				final int dimension = node.mDimension;

				if((this.mHigher == null || node.mPoint.getCoordinate(dimension).compareTo(this.mHigher.getCoordinate(dimension)) <= 0) && node.mHigher != null) {
					this.mKDNodeStack.addLast(node.mHigher);
				}

				if((this.mLower == null || node.mPoint.getCoordinate(dimension).compareTo(this.mLower.getCoordinate(dimension)) > 0) && node.mLower != null) {
					this.mKDNodeStack.addLast(node.mLower);
				}

				if(KDTree.this.isInRange(node.mPoint, this.mLower, this.mHigher)) {
					this.mNextKDNode = node;
					return old;
				}
			}

			this.mNextKDNode = null;

			return old;
		}

		@Override
		public void remove() throws UnsupportedOperationException {
			throw new UnsupportedOperationException();
		}
	}

	final class KeyIterator implements Iterator<P> {
		MapEntryIterator mMapEntryIterator;

		KeyIterator(final MapEntryIterator pMapEntryIterator) {
			this.mMapEntryIterator = pMapEntryIterator;
		}

		@Override
		public boolean hasNext() {
			return this.mMapEntryIterator.hasNext();
		}

		@Override
		public P next() {
			final Map.Entry<P, V> next = this.mMapEntryIterator.next();
			return (next == null ? null : next.getKey());
		}

		@Override
		public void remove() throws UnsupportedOperationException {
			this.mMapEntryIterator.remove();
		}
	}

	final class ValueIterator implements Iterator<V> {
		MapEntryIterator mMapEntryIterator;

		ValueIterator(final MapEntryIterator pMapEntryIterator) {
			this.mMapEntryIterator = pMapEntryIterator;
		}

		@Override
		public boolean hasNext() {
			return this.mMapEntryIterator.hasNext();
		}

		@Override
		public V next() {
			final Map.Entry<P, V> next = this.mMapEntryIterator.next();
			return (next == null ? null : next.getValue());
		}

		@Override
		public void remove() throws UnsupportedOperationException {
			this.mMapEntryIterator.remove();
		}
	}

	abstract class CollectionView<E> implements Collection<E> {

		@Override
		public boolean add(final E pItem) throws UnsupportedOperationException {
			throw new UnsupportedOperationException();
		}

		@Override
		public boolean addAll(final Collection<? extends E> pItems) throws UnsupportedOperationException {
			throw new UnsupportedOperationException();
		}

		@Override
		public void clear() {
			KDTree.this.clear();
		}

		@Override
		public boolean containsAll(final Collection<?> pItems) {
			for(final Object item : pItems) {
				if(!this.contains(item)) {
					return false;
				}
			}
			return true;
		}

		@Override
		public int hashCode() {
			return KDTree.this.hashCode();
		}

		@Override
		public boolean isEmpty() {
			return KDTree.this.isEmpty();
		}

		@Override
		public int size() {
			return KDTree.this.size();
		}

		@Override
		public Object[] toArray() {
			final Object[] array = new Object[this.size()];
			int i = 0;

			for(final E item : this) {
				array[i] = item;
				i++;
			}

			return array;
		}

		@Override
		@SuppressWarnings("unchecked")
		public <T> T[] toArray(T[] pArray) {
			Object[] array = pArray;

			if(array.length < this.size()) {
				array = pArray = (T[]) Array.newInstance(pArray.getClass().getComponentType(), this.size());
			}

			if(array.length > this.size()) {
				array[this.size()] = null;
			}

			int i = 0;
			for(final E item : this) {
				array[i] = item;
				i++;
			}

			return pArray;
		}
	}

	abstract class SetView<E> extends CollectionView<E> implements Set<E> {
		@Override
		public boolean equals(final Object pOther) {
			if(!(pOther instanceof Set<?>)) {
				return false;
			}

			if(pOther == this) {
				return true;
			}

			final Set<?> set = (Set<?>) pOther;

			if(set.size() != this.size()) {
				return false;
			}

			try {
				return this.containsAll(set);
			} catch (final ClassCastException cce) {
				return false;
			}
		}
	}

	final class MapEntrySet extends SetView<Map.Entry<P, V>> {
		@Override
		@SuppressWarnings("unchecked")
		public boolean contains(final Object pItem) throws ClassCastException, NullPointerException {
			final Map.Entry<P, V> e = (Map.Entry<P, V>) pItem;
			final KDNode node = KDTree.this.getNode(e.getKey());

			if(node == null) {
				return false;
			}

			return e.getValue().equals(node.getValue());
		}

		@Override
		public Iterator<Map.Entry<P, V>> iterator() {
			return new MapEntryIterator();
		}

		@Override
		@SuppressWarnings("unchecked")
		public boolean remove(final Object pItem) throws ClassCastException {
			final int size = this.size();
			final Map.Entry<P, V> e = (Map.Entry<P, V>) pItem;

			KDTree.this.remove(e.getKey());

			return (size != this.size());
		}

		@Override
		@SuppressWarnings("unchecked")
		public boolean removeAll(final Collection<?> pItems) throws ClassCastException {
			final int size = this.size();

			for(final Object item : pItems) {
				final Map.Entry<P, V> e = (Map.Entry<P, V>) item;
				KDTree.this.remove(e.getKey());
			}

			return (size != this.size());
		}

		@Override
		@SuppressWarnings("unchecked")
		public boolean retainAll(final Collection<?> pItems) throws ClassCastException {
			for(final Object item : pItems) {
				if(this.contains(item)) {
					final Collection<Map.Entry<P, V>> items = (Collection<Map.Entry<P, V>>) pItems;
					this.clear();
					for(final Map.Entry<P, V> e : items) {
						KDTree.this.put(e.getKey(), e.getValue());
					}
					return true;
				}
			}
			return false;
		}
	}

	final class KeySet extends SetView<P> {
		@Override
		public boolean contains(final Object pItem) throws ClassCastException, NullPointerException {
			return KDTree.this.containsKey(pItem);
		}

		@Override
		public Iterator<P> iterator() {
			return new KeyIterator(new MapEntryIterator());
		}

		@Override
		public boolean remove(final Object pItem) throws ClassCastException {
			final int size = this.size();
			KDTree.this.remove(pItem);
			return (size != this.size());
		}

		@Override
		public boolean removeAll(final Collection<?> pItems) throws ClassCastException {
			final int size = this.size();

			for(final Object o : pItems) {
				KDTree.this.remove(o);
			}

			return (size != this.size());
		}

		@Override
		@SuppressWarnings("unchecked")
		public boolean retainAll(final Collection<?> pItems) throws ClassCastException {
			final HashMap<P, V> map = new HashMap<P, V>();
			final int size = this.size();

			for(final Object item : pItems) {
				final V value = KDTree.this.get(item);

				if(value != null || this.contains(item)) {
					map.put((P) item, value);
				}
			}

			this.clear();
			KDTree.this.putAll(map);

			return (size != this.size());
		}
	}

	final class ValueCollection extends CollectionView<V> {
		@Override
		public boolean contains(final Object pItem) throws ClassCastException, NullPointerException {
			return KDTree.this.containsValue(pItem);
		}

		@Override
		public Iterator<V> iterator() {
			return new ValueIterator(new MapEntryIterator());
		}

		@Override
		public boolean remove(final Object pItem) throws ClassCastException {
			final KDNode node = KDTree.this.findValue(KDTree.this.mRoot, pItem);

			if(node != null) {
				KDTree.this.remove(node.getKey());
				return true;
			}

			return false;
		}

		@Override
		public boolean removeAll(final Collection<?> pItems) throws ClassCastException {
			final int size = this.size();

			for(final Object item : pItems) {
				KDNode node = KDTree.this.findValue(KDTree.this.mRoot, item);

				while(node != null) {
					KDTree.this.remove(item);
					node = KDTree.this.findValue(KDTree.this.mRoot, item);
				}
			}

			return (size != this.size());
		}

		@Override
		public boolean retainAll(final Collection<?> pItems) throws ClassCastException {
			final HashMap<P, V> map = new HashMap<P, V>();
			final int size = this.size();

			for(final Object item : pItems) {
				KDNode node = KDTree.this.findValue(KDTree.this.mRoot, item);

				while(node != null) {
					map.put(node.getKey(), node.getValue());
					node = KDTree.this.findValue(KDTree.this.mRoot, item);
				}
			}

			this.clear();
			KDTree.this.putAll(map);

			return (size != this.size());
		}
	}
}
